"""
基于神经网络的年径流预报实例
"""
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

np.set_printoptions(suppress=True, precision=4)
plt.style.use(['science', 'grid', 'muted'])


class DatasetLoader:
    """数据获取及预处理"""
    def __init__(self):
        # 载入数据
        data = np.loadtxt('data/jingliu.txt', delimiter='\t')
        # 数据预处理
        X, y = data[:, :-1], data[:, -1]
        X = MinMaxScaler().fit_transform(X)  # 规格化处理
        # 划分数据集
        self.X_train, self.X_test, self.Y_train, self.Y_test = \
            train_test_split(X, y, test_size=0.05, random_state=16)
        self.num_train_data, self.num_test_data = self.X_train.shape[
            0], self.X_test.shape[0]

    def get_batch(self, batch_size):
        # 从数据集中随机取出batch_size个元素并返回
        index = np.random.randint(0, self.num_train_data, batch_size)
        return self.X_train[index, :], self.Y_train[index]


# 构建网络
class MLP(tf.keras.Model):
    """多层感知机"""
    def __init__(self):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(units=64,
                                            activation=tf.nn.relu,
                                            input_shape=(4, ))
        self.dense2 = tf.keras.layers.Dense(units=128, activation=tf.nn.relu)
        self.dense3 = tf.keras.layers.Dense(units=64, activation=tf.nn.relu)
        self.dense4 = tf.keras.layers.Dense(units=64, activation=tf.nn.relu)
        self.dense5 = tf.keras.layers.Dense(units=1)

    def call(self, inputs):
        x = self.dense1(inputs)
        x = self.dense2(x)
        x = self.dense3(x)
        x = self.dense4(x)
        output = self.dense5(x)
        return output

    def hist_plot(self, history):
        plt.figure(figsize=(8, 6), dpi=128)
        plt.title('MLP Training Histotry')
        plt.xlabel('Epoch')
        plt.ylabel('Hist')
        plt.plot(history.epoch,
                 np.array(history.history['loss']),
                 label='Loss(MSE)')
        plt.plot(history.epoch,
                 np.array(history.history['mean_absolute_error']),
                 label='Metric(MAE)')
        plt.legend()
        plt.savefig('第15章：预测方法/MLP')


model = MLP()
data_loader = DatasetLoader()

num_epochs = 256
batch_size = 8
base_learning_rate = 1e-3

model.compile(
    optimizer=tf.keras.optimizers.RMSprop(learning_rate=base_learning_rate),
    loss=tf.keras.losses.mean_squared_error,
    metrics=[tf.keras.metrics.mean_absolute_error])
# 训练网络
history = model.fit(data_loader.X_train,
                    data_loader.Y_train,
                    epochs=num_epochs,
                    batch_size=batch_size)

print(model.summary())
model.hist_plot(history)

# 预测
y_pred = model.predict(data_loader.X_test).flatten()
print(data_loader.Y_test, y_pred)
